from typing import List
import re
from javalang import tokenizer
from .unicode import unicode_to_ascii

_NUM = {tokenizer.Integer,
        tokenizer.BinaryInteger,
        tokenizer.DecimalInteger,
        tokenizer.DecimalFloatingPoint,
        tokenizer.FloatingPoint,
        tokenizer.HexFloatingPoint,
        tokenizer.HexInteger,
        tokenizer.OctalInteger}

_LITERAL = {tokenizer.Character,
            tokenizer.Literal,
            tokenizer.String,
            tokenizer.Identifier}


def tokenize_code(code, subtoken=True, max_tokens=10000):
    """
    https://github.com/zy-zhou/HACS-release/blob/main/Data.py
    """
    code = unicode_to_ascii(code)
    tokens = list(tokenizer.tokenize(code))
    result = []
    for token in tokens:
        if type(token) in _NUM:
            result.append('num')
        elif type(token) in _LITERAL:
            value = node_split(token.value, subtoken)
            result.extend(value)
        else:
            result.append(token.value)
    return result[:max_tokens]


def node_split(s, subtoken=True):
    s = re.sub(r"\d+\.\d+\S*|0[box]\w*|\b\d+[lLfF]\b", " num ", s)
    s = re.sub(r"%\S*|[^A-Za-z0-9\s]", " ", s)
    s = re.sub(r"\b\d+\b", " num ", s)
    if subtoken:
        s = re.sub(r"[a-z][A-Z]", lambda x: x.group()[0] + " " + x.group()[1], s)
        s = re.sub(r"[A-Z]{2}[a-z]", lambda x: x.group()[0] + " " + x.group()[1:], s)
        s = re.sub(r"\w{32,}", " ", s)  # MD5, hash
        s = re.sub(r"[A-Za-z]\d+", lambda x: x.group()[0] + " ", s)
    s = re.sub(r"\s(num\s+){2,}", " num ", s)
    return s.lower().split()


def tokenize_doc(s: str, split_token: bool = True, min_tokens=0, max_tokens=10000) -> List[str]:
    """
    https://github.com/zy-zhou/HACS-release/blob/main/Data.py
    """
    s = unicode_to_ascii(s)
    if len(s) < min_tokens:
        return []
    s = re.sub(r"{@code|{@link", " ", s)
    s = re.sub(r"(([eE]\.[gG])|([iI]\.[eE]))\..+|<\S[^>]*>", " ", s)
    s = re.sub(r"\d+\.\d+\S*|0[box]\w*|\b\d+[lLfF]\b", " num ", s)
    first_p = re.search(r"[\.\?\!]+(\s|$)", s)
    if first_p is not None:
        s = s[:first_p.start()]
    s = re.sub(r"https:\S*|http:\S*|www\.\S*", " url ", s)
    s = re.sub(r"\b(todo|TODO)\b.*|[^A-Za-z0-9\.,\s]|\.{2,}", " ", s)
    s = re.sub(r"\b\d+\b", " num ", s)
    s = re.sub(r"([\.,]\s*)+", lambda x: " " + x.group()[0] + " ", s)

    if split_token:
        s = re.sub(r"[a-z][A-Z]", lambda x: x.group()[0] + " " + x.group()[1], s)
        s = re.sub(r"[A-Z]{2}[a-z]", lambda x: x.group()[0] + " " + x.group()[1:], s)
        s = re.sub(r"\w{32,}", " ", s)  # MD5
        s = re.sub(r"[A-Za-z]\d+", lambda x: x.group()[0] + " ", s)
    s = re.sub(r"\s(num\s+){2,}", " num ", s)
    s = s.lower().split()
    return s[:max_tokens]
